import numpy as np
import cv2
import os
import post_process as post

# import sys
# sys.path.append('/home/solan/cxk/caffe/python')

import caffe

# resnet+fpn
# net_file= 'model/res-20class-41.7mAP/retinanet.prototxt'
# caffe_model='model/res-20class-41.7mAP/retinanet.caffemodel'
# mobilenet v2+fpn
net_file= 'model/retinanet.prototxt'
caffe_model='model/retinanet.caffemodel'
# net_file= 'no_bn.prototxt'
# caffe_model='no_bn.caffemodel'
net = caffe.Net(net_file,caffe_model,caffe.TEST)

test_dir = "images"
CLASSES = ('person', 'aeroplane', 'tv', 'train',
           'boat', 'dog', 'chair', 'bird', 'bicycle',
           'bottle', 'sheep', 'dinigtable', 'horse',
           'motorbike', 'sofa', 'cow',
           'car', 'cat', 'bus', 'pottedplant')


def preprocess(src, size):

    image = cv2.resize(src, size)

    image = image.astype(np.float32)
    image /= 255
    image -= [0.485, 0.456, 0.406]
    image /= [0.229, 0.224, 0.225] # (608, 608, 3)
    image = np.transpose(image, (2, 0, 1)) # (3, 608, 608)
    return image

def postprocess(img, out):
    h = img.shape[0]
    w = img.shape[1]
    box = out['detection_out'][0,0,:,3:7] * np.array([w, h, w, h])

    cls = out['detection_out'][0,0,:,1]
    conf = out['detection_out'][0,0,:,2]
    return (box.astype(np.int32), conf, cls)

def detect(imgfile):
    input_size = 608
    origimg = cv2.imread(imgfile)
    img = preprocess(origimg, (input_size, input_size)) # (3, 608, 608)

    net.blobs['blob1'].data[...] = img
    out = net.forward() # dict (copy content)

    # resnet50+fpn
    # reg_label = ['conv_blob66', 'conv_blob71', 'conv_blob76', 'conv_blob81', 'conv_blob86']
    # class_label = ['conv_blob91', 'conv_blob96', 'conv_blob101', 'conv_blob106', 'conv_blob111']
    # mobilenet v2+fpn
    reg_label = ['conv_blob64', 'conv_blob69', 'conv_blob74', 'conv_blob79', 'conv_blob84']
    class_label = ['conv_blob89', 'conv_blob94', 'conv_blob99', 'conv_blob104', 'conv_blob109']

    num_classes = 20
    # anchor_num = 9 # anchor num of a pix in feature

    for key, value in out.items():
        out[key] = np.transpose(value, (0, 2, 3, 1))
        if key in reg_label:
            out[key] = out[key].reshape(-1, 4)
        elif key in class_label:
            out[key] = out[key].reshape(-1, num_classes)
        # print (key, out[key].shape)

    regression = np.concatenate((out[reg_label[0]],out[reg_label[1]],out[reg_label[2]],
                                 out[reg_label[3]],out[reg_label[4]]), 0) # [x, 4]

    anchors = post.get_anchor() # [x, 4] (x1 y1 x2 y2)
    transformed_anchors = post.regression(anchors, regression) # get bbox shape[x, 4] (x1 y1 x2 y2)
    transformed_anchors = post.clipBoxes(transformed_anchors, img)

    classification = np.concatenate((out[class_label[0]], out[class_label[1]], out[class_label[2]],
                                     out[class_label[3]], out[class_label[4]]), 0)  # [x, 20]

    scores = classification.max(axis=1)
    scores = 1 / (1 + np.exp(-1*scores))
    classes = classification.argmax(axis=1)
    # scores = scores.unsqueeze(1).unsqueeze(0)
    # scores = torch.max(classification, dim=2, keepdim=True)[0] # [1, 111330, 1]
    scores_over_thresh = (scores > 0.5)  #
    if scores_over_thresh.sum() == 0: # no boxes to NMS
        return True

    classification = classification[scores_over_thresh, :] # [x]
    transformed_anchors = transformed_anchors[scores_over_thresh, :] # [x, 4]
    scores = scores[scores_over_thresh] # [x]
    classes = classes[scores_over_thresh] # [x]


    nms_scores, nms_class, nms_bboxes = post.CIoU_nms(scores[:], classes, transformed_anchors[:, :],
                                                      nms_threshold=0.5, flag=0)  # flag[0-CIoU; 1-DIoU; 2-IoU]

    for i in range(len(nms_bboxes)):
        nms_bboxes[i][::2] = nms_bboxes[i][::2] * origimg.shape[1] / input_size
        nms_bboxes[i][1::2] = nms_bboxes[i][1::2] * origimg.shape[0] / input_size
        p1 = (int(nms_bboxes[i][0]), int(nms_bboxes[i][1]))
        p2 = (int(nms_bboxes[i][2]), int(nms_bboxes[i][3]))
        cv2.rectangle(origimg, p1, p2, (0,255,0))
        p3 = (max(p1[0], 15), max(p1[1], 15))
        title = "%s:%.2f" % (CLASSES[int(nms_class[i])], nms_scores[i])
        cv2.putText(origimg, title, p3, cv2.FONT_ITALIC, 0.6, (0, 255, 0), 1)
        cv2.imshow("SSD", origimg)

    k = cv2.waitKey(0) & 0xff
        #Exit if ESC pressed
    if k == 27 : return False
    return True

for f in os.listdir(test_dir):
    if detect(test_dir + "/" + f) == False:
       break


